import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
from matplotlib import pyplot as plt


def getTransform():
    """
    @brief return transform that covert picture into tensor
    @return transform object
    """
    # 设置transform 对图像进行预处理
    transform = transforms.Compose(
        [
            # 将Pic转换为Tensor 并转化到[0,1]
            transforms.ToTensor(),
            # 对输入标准化 均值0.5 标准差0.5
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ]
    )

    return transform


def getDataSet(root_path: str, batch_size: int, transform):
    """
    @brief return dataset build by torchvision
    @param root_path: where will data be stored
    @param batch_size: size of each batch
    @param transform: what transform will be applied to input data
    """

    # 训练数据集
    trainset = torchvision.datasets.CIFAR10(
        root=root_path,  # 数据根目录
        train=True,  # 是训练集
        download=True,  # 自动下载
        transform=transform,  # 预处理方法
    )

    # 训练数据加载器
    trainloader = torch.utils.data.DataLoader(
        trainset,  # 设置数据集
        batch_size=batch_size,  # 分批大小
        shuffle=True,  # 打乱顺序 是
        num_workers=0,  # 0 workers工作
        pin_memory=True if torch.cuda.is_available else False
    )

    # 测试数据集
    testset = torchvision.datasets.CIFAR10(
        root=root_path,
        train=False,  # 是测试集
        download=True,
        transform=transform
    )

    # 测试数据加载器
    testloader = torch.utils.data.DataLoader(
        testset,
        batch_size=batch_size,
        shuffle=False,  # 打乱顺序 否
        num_workers=0,
        pin_memory=True if torch.cuda.is_available else False
    )

    return trainset, trainloader, testset, testloader


def imshow(img):
    """
    @brief show img which have been normalized and convert into tensor
    @param img: img tensor
    """
    img = img / 2 + 0.5  # 去标准化
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


class DataPrinter:
    """
    @brief Using for plot loss and accuracy curve
    """

    def __init__(self):
        self.acc = []
        self.loss = []
        self.fig = plt.figure()

    def __call__(self, loss, acc):
        """
        @brief Add data call
        @param loss: loss val
        @param acc: accuracy val
        """
        self.acc.append(acc)
        self.loss.append(loss)

    def show(self):
        """
        @brief call for show data
        """
        x = list(range(1, len(self.acc) + 1))
        self.fig.clf()
        ax1 = self.fig.add_subplot(111)
        ax1.plot(x, self.acc, color='red', label='test_accuracy')
        ax1.set_xlabel('epoch')
        ax1.set_ylabel('test_accuracy')
        ax2 = ax1.twinx()
        ax2.plot(x, self.loss, color='blue', label='train_loss')
        ax2.set_ylabel('train_loss')
        self.fig.legend(loc='upper right', bbox_to_anchor=(1, 1), bbox_transform=ax1.transAxes)
        self.fig.show()

    def clear(self):
        """
        @brief clear figure and data
        """
        self.loss.clear()
        self.acc.clear()
        self.fig.clf()

    def save(self, path: str):
        """
        @brief save figure to file
        @param path: save path
        """
        self.fig.savefig(path)
